import torch
import torch.fft
import torch.nn.functional as F

class FeatureFFTStyleTransfer:
    def __init__(self, radius=0.1, beta=1.0):
        self.radius = radius
        self.beta = beta

    def fft2d(self, x):
        return torch.fft.fftshift(torch.fft.fft2(x, norm='ortho'), dim=(-2, -1))

    def ifft2d(self, x):
        return torch.fft.ifft2(torch.fft.ifftshift(x, dim=(-2, -1)), norm='ortho').real

    def low_high_split(self, amp):
        B, C, H, W = amp.shape
        cy, cx = H // 2, W // 2
        y = torch.arange(H).to(amp.device).reshape(1, 1, -1, 1)
        x = torch.arange(W).to(amp.device).reshape(1, 1, 1, -1)
        dist = torch.sqrt((y - cy)**2 + (x - cx)**2)
        mask = (dist < (self.radius * max(H, W))).float()
        return amp * mask, amp * (1 - mask)

    def transfer(self, feat_src, feat_tgt, align=False):
        """
        feat_src, feat_tgt: [B, C, H, W]
        """
        fft_src = self.fft2d(feat_src)
        fft_tgt = self.fft2d(feat_tgt)

        amp_src, phase_src = torch.abs(fft_src), torch.angle(fft_src)
        amp_tgt = torch.abs(fft_tgt)

        amp_src_low, amp_src_high = self.low_high_split(amp_src)
        amp_tgt_low, _ = self.low_high_split(amp_tgt)

        # amp_fused = amp_src_high + self.beta * amp_tgt_low
        # fused = amp_fused * torch.exp(1j * phase_src)

        # feat_fused = self.ifft2d(fused)
        # if align:
        # 高频部分解耦出来做语义对齐
        high_src = amp_src_high * torch.exp(1j * phase_src)
        high_tgt = amp_src_high * torch.exp(1j * torch.angle(self.fft2d(feat_tgt)))  # 保证结构一致
        feat_src_high = self.ifft2d(high_src)
        feat_tgt_high = self.ifft2d(high_tgt)
        loss_sem = F.mse_loss(feat_src_high, feat_tgt_high)
        return loss_sem
        # return feat_fused
